import vosk
import pyaudio
import time
from gtts import gTTS
import os

# 关键词，当识别到这个关键词时，程序会作出回应
KEYWORD = "小爱同学"

class XiaoAiAssistant:
    def __init__(self):
        self.model_path = "E:\\work\\project_st\\vosk-model-small-cn-0.22"  # 替换为你的模型文件路径
        self.model = vosk.Model(self.model_path)
        self.p = pyaudio.PyAudio()
        self.recognizer = None
        self.stream = self.p.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True, frames_per_buffer=4096)

    def speak(self, text):
        tts = gTTS(text=text, lang='zh-cn')
        tts.save('temp.mp3')
        self.play_audio('temp.mp3')

    def play_audio(self, filename):
        # 播放音频文件
        with open(filename, 'rb') as audio_file:
            data = audio_file.read()
        self.stream.stop_stream()
        self.stream.close()
        play_stream = self.p.open(format=pyaudio.paInt16, channels=1, rate=16000, output=True)
        play_stream.write(data)
        play_stream.close()
        self.stream = self.p.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True, frames_per_buffer=4096)
        self.stream.start_stream()

    def listen(self):
        print(f"监听中，说'{KEYWORD}'来唤醒我。")
        self.recognizer = vosk.KaldiRecognizer(self.model, 16000)
        while True:
            data = self.stream.read(4096)
            if self.recognizer.AcceptWaveform(data):
                result = self.recognizer.Result()
                if result:
                    result = result.replace(" ", "")
                    print(result)
                
                
                # if self.recognizer.FinalResult():
                #     result = self.recognizer.Result()
                #     print(result)
                    if KEYWORD in result:
                        self.speak("我在这里")
                else:
                    result = self.recognizer.PartialResult()
                    print("部分识别结果:", result)

    def start(self):
        self.recognizer = vosk.KaldiRecognizer(self.model, self.p.get_device_info_by_index(0).get('defaultSampleRate'))
        self.stream.start_stream()
        self.listen()

if __name__ == "__main__":
    assistant = XiaoAiAssistant()
    assistant.start()